#ifndef SILL_LEARNING_DMOVE_PUSH_BACK_SUBTREE_HPP
#define SILL_LEARNING_DMOVE_PUSH_BACK_SUBTREE_HPP

#include <sill/learning/structure_old/decomposable_move.hpp>

#include <sill/macros_def.hpp>

namespace sill {

  /**
   * Class for representing a possible move in structure search for
   * decomposable models:
   *  One variable is replaced with another in a single
   *  node.  This may be thought of as the subtree induced by one variable
   *  pushing back the subtree induced by another variable (though this is
   *  a bit misleading since the subtrees may overlap).
   *
   * For maximal and non-maximal JTs.
   *
   * @param F         type of factor used in the decomposable model
   * @param Inserter  Inserter.push(m, score) must insert move m with
   *                  score/estimate score into whatever container collects
   *                  new moves which are generated by this class.
   *
   * \author Joseph Bradley
   * \ingroup learning_structure
   */
  template <typename F, typename Inserter>
  class dmove_push_back_subtree
    : public decomposable_move<F,Inserter> {

  public:
    //! The base (move) type
    typedef decomposable_move<F,Inserter> base;

    //! The type of factor used in the decomposable model
    typedef F factor_type;

    //! The type of variable associated with a factor
    typedef typename F::variable_type variable_type;

    //! The domain type of the factor
    typedef typename base::domain_type domain_type;

    //! The generic representation for clique changes
    typedef typename base::clique_change clique_change;

    //! The type of edge associated with the model
    typedef typename base::edge edge;

    //! The type of vertex associated with the model
    typedef typename base::vertex vertex;

    //////////////////////// PRIVATE DATA AND METHODS ///////////////////////

  private:
    //! Vertex being modified
    vertex v;

    //! Variable to add to vertex v
    variable_type* add_var;

    //! Variable to remove from vertex v
    variable_type* remove_var;

    //! Find all variables which appear in u's neighbors exactly once.
    domain_type get_once_vars(const vertex& u,
                              const decomposable<F>& model) const {
      std::map<variable_type*, size_t> var_count;
      foreach(const vertex& u2, model.neighbors(u)) {
        foreach(variable_type* tmpvar, model.clique(u2))
          var_count[tmpvar]++;
      }
      domain_type once_vars;
      typename std::map<variable_type*, size_t>::const_iterator it_end
        = var_count.end();
      for (typename std::map<variable_type*, size_t>::const_iterator it
             = var_count.begin();
           it != it_end; ++it)
        if (it->second == 1)
          once_vars.insert(it->first);
      return once_vars;
    }

    //! Generate moves (remove r, add a) for r in r_vars and a in a_vars
    //! and vertex u.
    void
    generate_moves_(const vertex& u, const domain_type& r_vars,
                    const domain_type& a_vars,
                    const learnt_decomposable<F>& model, double cur_score,
                    const decomposable_score<F>& score,
                    dataset_statistics<>& stats, Inserter& inserter,
                    bool use_estimates) const {
      foreach(variable_type* r, r_vars) {
        foreach(variable_type* a, a_vars) {
          dmove_push_back_subtree<F,Inserter>* move_ptr =
            new dmove_push_back_subtree<F,Inserter>(u, a, r);
          double change;
          if (use_estimates)
            change =
              score.estimate_change(model, cur_score, *move_ptr, stats).second;
          else
            change
              = score.compute_change(model, cur_score, *move_ptr, stats).second;
          inserter.push(move_ptr, change);
        }
      }
    }

    ////////////////////////// PUBLIC METHODS //////////////////////////////
  public:
    //! Creates a null move (which can be used to generate new moves).
    dmove_push_back_subtree() : v(decomposable<F>::null_vertex()) { }

    //! Creates a move which replaces remove_var with add_var in vertex v.
    dmove_push_back_subtree(vertex v, variable_type* add_var,
                           variable_type* remove_var)
      : v(v), add_var(add_var), remove_var(remove_var) {
    }

    //! Given a model, score, and statistics class, generate all possible moves
    //! and insert pointers to them (and their scores) into the provided
    //! Inserter.
    void
    generate_all_moves(const learnt_decomposable<F>& model, double cur_score,
                       const decomposable_score<F>& score, dataset_statistics<>& stats,
                       Inserter& inserter, bool use_estimates = false) const {
      // Foreach vertex
      foreach(const vertex& u, model.vertices()) {
        // Find variable pairs (a,r) where each is in exactly one
        //  of u's neighbors but only r is in u.
        // First, find all variables which appear in u's neighbors exactly once.
        domain_type once_vars(get_once_vars(u, model));
        // Now, divide once_vars into those which do/don't appear in u.
        domain_type not_in_u(set_difference(once_vars, model.clique(u)));
        once_vars = set_difference(once_vars, not_in_u);
        // Generate moves.
        generate_moves_(u, once_vars, not_in_u, model, cur_score, score,
                        stats, inserter, use_estimates);
      }
    }

    //! Given a model, score, another decomposable_move (which has just been
    //! committed), and statistics class,
    //! generate all possible new moves of this type and insert pointers to
    //! them (and their scores) into the provided Inserter.
    void
    generate_new_moves(const learnt_decomposable<F>& model, double cur_score,
                       const decomposable_score<F>& score,
                       const std::vector<clique_change>& clique_changes,
                       dataset_statistics<>& stats, Inserter& inserter,
                       bool use_estimates = false) const {
      // For each changed clique,
      foreach(const clique_change& change, clique_changes) {
        domain_type once_vars(get_once_vars(change.v, model));
        if (change.deleted_vars.size() > 0) {
          // Variables {a} were deleted, so
          // check for moves (add a, remove r) for this clique.
          // (Find a in adjacent cliques, and find variables
          //  r in this clique and exactly one adjacent clique.)
          domain_type r_vars(set_intersect(once_vars, model.clique(change.v)));
          domain_type a_vars(set_intersect(change.deleted_vars, once_vars));
          generate_moves_(change.v, a_vars, r_vars, model, cur_score, score,
                          stats, inserter, use_estimates);
        }
        if (change.added_vars.size() > 0) {
          // Variables {u} were added, so
          // check for moves (add u2, remove u) for this clique.
          // (Find u in exactly one adjacent clique, and find u2 not in this
          //  clique and in an adjacent clique.)
          domain_type r_vars(set_intersect(change.added_vars, once_vars));
          domain_type a_vars(set_difference(once_vars, model.clique(change.v)));
          generate_moves_(change.v, a_vars, r_vars, model, cur_score, score,
                          stats, inserter, use_estimates);
          // Check for moves (add u, remove u2) for adjacent cliques.
          // (Find neighbors v2 without u and with variables u2 where u2 appears
          //  in exactly one neighbor of v2.)
          foreach(const vertex& v2, model.neighbors(change.v)) {
            const domain_type& v2_clique = model.clique(v2);
            a_vars = set_difference(change.added_vars, v2_clique);
            domain_type v2_once_vars(get_once_vars(v2, model));
            r_vars = set_intersect(v2_clique, v2_once_vars);
            generate_moves_(v2, a_vars, r_vars, model, cur_score, score,
                            stats, inserter, use_estimates);
          }
        }
      }
    }

    //! Call a score functor on each clique/separator inserted/deleted
    //! by this move (where the move is represented by clique/separator
    //! insertions/deletions).
    //! @return false if move is invalid, else true
    bool map_score_functor(decomposable_score_functor<F>& func,
                           const learnt_decomposable<F>& model,
                           dataset_statistics<>& stats) const {
      if (!(model.contains(v)) || model.clique(v).count(add_var)
          || !(model.clique(v).count(remove_var)))
        return false;
      func.deleted_clique(model.marginal(v));
      func.inserted_clique(stats.marginal(set_union(set_difference(model.clique(v), remove_var), add_var)));
      size_t n_add = 0; // number of edges containing add_var
      size_t n_remove = 0; // number of edges containing remove_var
      foreach(const edge& e, model.out_edges(v)) {
        const domain_type& s = model.separator(e);
        if (model.clique(e.target()).count(add_var)) {
          ++n_add;
          func.deleted_separator(model.marginal(e));
          if (s.count(remove_var)) {
            ++n_remove;
            func.inserted_separator(stats.marginal(set_union(set_difference(s, remove_var), add_var)));
          } else {
            func.inserted_separator(stats.marginal(set_union(s, add_var)));
          }
        } else if (s.count(remove_var)) {
          ++n_remove;
          func.deleted_separator(model.marginal(e));
          const domain_type& s = model.separator(e);
          func.inserted_separator(stats.marginal(set_union(set_difference(s, remove_var), add_var)));
        }
      }
      if (n_add != 1 || n_remove != 1)
        return false;
      return true;
    }

    //! Given a model, check to see if the move is still valid.
    bool valid(const learnt_decomposable<F>& model) const {
      if (v == decomposable<F>::null_vertex() || !(model.contains(v)))
        return false;
      const domain_type& c = model.clique(v);
      if (c.count(add_var) || !(c.count(remove_var)))
        return false;
      // Make sure that add_var, remove_var are each contained in exactly
      //  one neighbor.
      size_t n_add = 0; // number of edges containing add_var
      size_t n_remove = 0; // number of edges containing remove_var
      foreach(const vertex& v2, model.neighbors(v)) {
        const domain_type& c2 = model.clique(v2);
        if (c2.count(add_var))
          ++n_add;
        if (c2.count(remove_var))
          ++n_remove;
      }
      if (n_add != 1 || n_remove != 1)
        return false;
      return true;
    }

    //! Given a mutable model, commit the move.
    //! Note: This does NOT check if the move is valid!
    //! Also, this does not calibrate or renormalize the model.
    std::vector<clique_change>
    commit(learnt_decomposable<F>& model, dataset_statistics<>& stats) const {
      if (v == decomposable<F>::null_vertex())
        assert(false);
      domain_type d(model.clique(v));
      d.insert(add_var);
      d.erase(remove_var);
      model.set_clique(v, d, stats.marginal(d));
      domain_type removedomain; removedomain.insert(remove_var);
      domain_type adddomain; adddomain.insert(add_var);
      return std::vector<clique_change>(1, clique_change(v,removedomain,adddomain));
    }

  }; // class dmove_push_back_subtree

} // namespace sill

#include <sill/macros_undef.hpp>

#endif // #ifndef SILL_LEARNING_DMOVE_PUSH_BACK_SUBTREE_HPP
